import os
import torch
from PIL import Image
from torchvision.utils import save_image
from tqdm import tqdm
from segmentation.utils import evaluate_bce, evaluate_ce, visualize, get_save_image
from torchvision import transforms


def apply_trans(input):
    trans = transforms.Compose([
        transforms.RandomHorizontalFlip(0.5),
        transforms.RandomAffine(degrees=10, translate=(0.25,0.25), scale=(0.75,1.25), resample=Image.NEAREST, fillcolor=None)
        # transforms.RandomAffine(degrees=0, translate=(0.1,0.1), scale=(0.75,1), resample=Image.NEAREST, fillcolor=None)
    ])

    return trans(input)


def apply_trans_batch(image, mask, chn):
    B = image.shape[0]
    image = (image+1)/2
    for i in range(B):
        image_, mask_ = apply_trans(torch.cat([image[i], mask[i].unsqueeze(0)], dim=0)).split([chn, 1], dim=0)
        image[i], mask[i] = image_, mask_.squeeze(0).to(torch.long)

    image = (image-0.5)/0.5

    return image, mask


def eval_net_supervised(net, loader, is_test, folder_dir, i_train, multi_class, ignore_bg=False, bg_label=0):
    net.eval()
    evaluate = evaluate_ce if multi_class else evaluate_bce
    sum_dice, sum_accuracy, sum_iou, sum_wdice, sum_wacc, sum_wiou = 0,0,0,0,0,0
    to_save = []
    count = 0
    nSample = 0
    for _, image, mask in tqdm(loader, total=len(loader), desc='Validation', unit='batch', leave=False):
        nSample += image.shape[0]
        image, mask = image.cuda(), mask.cuda()
        image = (image-0.5)/0.5

        pred, _ = net(image)
        
        if multi_class:
            dice, accuracy, iou, wdice, wacc, wiou = evaluate(pred, mask, parts=pred.shape[1], ignore_bg=ignore_bg, bg_label=bg_label)
            sum_wdice    += wdice
            sum_wacc     += wacc
            sum_wiou     += wiou
        else:
            dice, accuracy, iou = evaluate(pred, mask, ignore_bg=ignore_bg, bg_label=bg_label)

        sum_dice     += dice
        sum_accuracy += accuracy
        sum_iou      += iou

        if is_test:
            image = (image+1)/2

            if multi_class:
                # B, 1, H, W -> B, 3, H, W
                image = image.repeat(1,3,1,1)
                # B, 5, H, W -> B, 3, H, W
                pred = visualize(pred.argmax(1))
                mask = visualize(mask)
                
                to_save.append(torch.cat([
                    torch.cat([image.unsqueeze(1), mask.unsqueeze(1)], dim=1).view(-1, 3, 256, 256),
                    torch.cat([image.unsqueeze(1), pred.unsqueeze(1)], dim=1).view(-1, 3, 256, 256)
                ], dim=0))
            else:
                pred = torch.where(pred>0, 1.0, 0.0)
                to_save.append(get_save_image(image.detach().cpu(), mask.detach().cpu(), image.detach().cpu(), pred.detach().cpu()))

            if len(to_save) == 10:
                save_image(torch.cat(to_save, dim=0), os.path.join(folder_dir, 'visualization', f'test{i_train}_{count}.png'), nrow=10)
                to_save = []
                count += 1

    if is_test and len(to_save)>0:
        save_image(torch.cat(to_save, dim=0), os.path.join(folder_dir, 'visualization', f'test{i_train}_{count}.png'), nrow=10)

    return sum_dice/nSample, sum_accuracy/nSample, sum_iou/nSample, sum_wdice/nSample, sum_wacc/nSample, sum_wiou/nSample
